import json
import re
from collections import defaultdict
from huey.exceptions import TaskException

from src.tools.chain import gen_chain
from src.tools.extract_jsonstring import re_json_string, re_json_string_list
from src.gen_file.prompt import SOPPrompt
from src.tools.dms_operator import dmsoperator
from src.tools.information_exchange import pub_gen
from src.tools.split_string import split_sub_string
from src.tools.mml_to_latex import mathml2latex
from src.tools.file_type import check_file_type1

class GenSOPFile:

    def __init__(self):
        self.prompt = SOPPrompt()

    # 解析工艺规程文件，提取出需要修改的段落和表格
    def parse_SOPjsonfile(self, file: dict):
        #获得文件对应字典数据
        json_data = dmsoperator.get_jsonfile(file.get("realPath"))
        # 构造修改的新数据
        file_pubinfo = {
            "filePath": json_data["filePath"],
            "isHighlight": True,
            "editInfos": []
        }
        # 提取需要修改的段落
        parainfo = []
        for para in json_data["paragraphInfos"]:
            is_highlightcolor_para = False

            para_text = ""
            # 提取纯文字和下标
            for textinfo in para["textInfos"]:
                if textinfo["subSuperscript"] == 2:
                    para_text += f'<sub>{textinfo["text"]}</sub>'
                else:
                    para_text += textinfo["text"]
                if textinfo["highlightColor"] != "0, 0, 0, 0":
                    is_highlightcolor_para = True
            # 提取公式
            mathinfos = []
            for mathinfo in para["officeMathInfos"]:
                if mathinfo["highlightColor"] != "0, 0, 0, 0":
                    is_highlightcolor_para = True
                math_dict = {
                    "text": mathml2latex(mathinfo["mathMlCode"])
                }
                mathinfos.append(math_dict)

            para_dict = {
                "index": para["index"],
                "sectionIndex": para["sectionIndex"],
                "bodyChildIndex": para["bodyChildIndex"],
                "text": para_text,
                "latexMathInfos": mathinfos
            }
            if is_highlightcolor_para:
                parainfo.append(para_dict)
        # 提取需要修改的表格
        tableinfo = []
        for table in json_data["tableInfos"]:
            # 标注高亮是修改部分, 以table为单位进行修改
            is_highlightcolor_table = False
            table_dict = {
                "index": table["index"],
                "sectionIndex": table["sectionIndex"],
                "bodyChildIndex": table["bodyChildIndex"],
                "tableCellInfos": []
            }
            for cell in table["tableCellInfos"]:
                cell_dict = {
                    "rowIndex": cell["rowIndex"],
                    "columnIndex": cell["columnIndex"],
                    "paragraphInfos": []
                }
                for para in cell["paragraphInfos"]:
                    para_text = ""
                    for textinfo in para["textInfos"]:
                        if textinfo["subSuperscript"] == 2:
                            para_text += f'<sub>{textinfo["text"]}</sub>'
                        else:
                            para_text += textinfo["text"]
                        if textinfo["highlightColor"] != "0, 0, 0, 0":
                            is_highlightcolor_table = True

                    mathinfos = []
                    for mathinfo in para["officeMathInfos"]:
                        if mathinfo["highlightColor"] != "0, 0, 0, 0":
                            is_highlightcolor_table = True
                        math_dict = {
                            "text": mathml2latex(mathinfo["mathMlCode"])
                        }
                        mathinfos.append(math_dict)

                    para_dict = {
                        "index": para["index"],
                        "sectionIndex": para["sectionIndex"],
                        "bodyChildIndex": para["bodyChildIndex"],
                        "text": para_text,
                        "latexMathInfos": mathinfos
                    }
                    cell_dict["paragraphInfos"].append(para_dict)
                table_dict["tableCellInfos"].append(cell_dict)
            if is_highlightcolor_table:
                tableinfo.append(table_dict)
        return file_pubinfo, parainfo, tableinfo

    # 找到流式输出内容
    def find_key_values(self, target_key, data):
        found_values = []
        if isinstance(data, dict):
            for key, value in data.items():
                if key == target_key and value:
                    found_values.append(value)
                found_values.extend(self.find_key_values(target_key, value))
        elif isinstance(data, list):
            for item in data:
                found_values.extend(self.find_key_values(target_key, item))
        return found_values

    # 利用生成链，生成要修改的内容
    def gen_SOPfile_para_table(self, content, extract_text, task_id, SOPfile):
        # 修改对应内容
        revise_SOPtext_query = self.prompt.gen_file_prompt(orig_SOPtext=json.dumps(content, ensure_ascii=False),
                                                          pitext=extract_text)
        # print("-" * 100)
        # print(revise_SOPtext_query)
        # print("-" * 100)
        revise_SOPtext = gen_chain(revise_SOPtext_query)
        try:
            revise_SOPtext = re_json_string(revise_SOPtext)
            print(revise_SOPtext)
            revise_SOPtext_dict = json.loads(revise_SOPtext)
        except:
            print("json解析失败")
            revise_SOPtext_dict = content
        # 找到流式输出内容
        text_key = self.find_key_values("text", revise_SOPtext_dict)
        for msg in text_key:
            print(f"生成send: {msg}")
            pub_gen(task_id, msg, SOPfile.get("docId"), SOPfile.get("taskDocId"))
        return revise_SOPtext_dict

    # 有来源的修改
    def gen_SOPfile_para_table_temp(self, content, extract_text, template_list, task_id, SOPfile):
        template_text = []
        for temp in template_list:
            template_text.append({"text": temp["text"], "tableCellInfos": temp["tableCellInfos"]})
        revise_SOPtext_query = self.prompt.gen_file_prompt_temp(orig_SOPtext=json.dumps(content, ensure_ascii=False),
                                                                pitext=extract_text, temp=str(template_text))
        print("-" * 100)
        print(revise_SOPtext_query)
        print("-" * 100)
        revise_SOPtext = gen_chain(revise_SOPtext_query)
        revise_SOPtext = re_json_string(revise_SOPtext)
        print(revise_SOPtext)
        revise_SOPtext_dict = json.loads(revise_SOPtext)
        text_key = self.find_key_values("text", revise_SOPtext_dict)
        for msg in text_key:
            print(f"send11111{msg}")
            pub_gen(task_id, msg, SOPfile.get("docId"), SOPfile.get("taskDocId"))
        return revise_SOPtext_dict

    # 正则判断是否以附件数字开头
    def starts_with_attachment_number(self, text):
        # 定义正则表达式，确保仅匹配“附件”开头，后跟正整数
        pattern = r"^附件[1-9]\d*(?!\.\d)"
        # 使用 re.match 检查是否匹配
        return bool(re.match(pattern, text))

    # 去掉不重要的信息
    def cut_filetext(self, filepath):
        pifile_json = dmsoperator.get_jsonfile(filepath)
        combined = pifile_json["paragraphInfos"] + pifile_json["tableInfos"]
        # 按照 index 排序重新组合json文件
        sorted_combined_pifile = sorted(combined, key=lambda x: x["index"])
        # 通过正则拿到附件标题
        title_pos = defaultdict(list)
        for para in pifile_json["paragraphInfos"]:
            if self.starts_with_attachment_number(para["text"]):
                title_pos[para["text"]].append({"index": para["index"], "sectionIndex": para["sectionIndex"], "bodyChildIndex": para["bodyChildIndex"]})
        titles = [title for title in title_pos.keys()]
        # 找到要去掉的标题
        titles_text = json.dumps(titles, ensure_ascii=False)
        cut_prompt = self.prompt.cut_prompt(titles_text)
        un_titles_text = gen_chain(cut_prompt)
        print(un_titles_text)
        un_titles = json.loads(re_json_string_list(un_titles_text))
        print(un_titles)
        # 去掉标题对应所有内容
        need_titles = []
        title_index = 0
        for i, _ in enumerate(titles):
            if titles[i] == un_titles[title_index]:
                if i < len(titles) - 1:
                    need_titles.append([titles[i], titles[i + 1]])
                    title_index += 1
                    if title_index == len(un_titles):
                        break
                else:
                    need_titles.append([titles[i], ""])
        print(need_titles)
        dele_fileinfo = {
            "filePath": filepath,
            "editInfos": []
        }
        for strat_title, end_title in need_titles:
            for info in sorted_combined_pifile:
                if end_title != "" and title_pos[strat_title][-1]["index"] <= info["index"] < title_pos[end_title][-1]["index"]:
                    dele_fileinfo["editInfos"].append(
                        {
                            "index": info["index"],
                            "sectionIndex": info["sectionIndex"],
                            "bodyChildIndex": info["bodyChildIndex"],
                            "operationType": 2
                        }
                    )
                if end_title == "" and title_pos[strat_title][-1]["index"] <= info["index"]:
                    dele_fileinfo["editInfos"].append(
                        {
                            "index": info["index"],
                            "sectionIndex": info["sectionIndex"],
                            "bodyChildIndex": info["bodyChildIndex"],
                            "operationType": 2
                        }
                    )

        new_filepath = dmsoperator.post_edit_file_v2(dele_fileinfo)["data"]

        need_text = dmsoperator.get_fileconversion(new_filepath, 4).replace(
            'style="text-align: center;"', "").replace('style="text-align: left;"', "").replace(
            'style="text-align: right;"', "")
        return need_text


    # 生成工艺规程文件
    def gen_SOPfile(self, pi_query: str, pifiles: list, SOPfile: dict, template_infos: list, task_id: str, change_info: dict):
        extract_text, pifile_id = "", ""
        # 判断内容来源文件是否有工艺信息表
        for pifile in pifiles:
            pifile_id_temp = pifile.get("realPath")
            if "工艺信息表" in pifile.get("documentName") or "工艺信息表" in check_file_type1(pifile_id_temp):
                pifile_id = pifile_id_temp
        # 没有工艺信息表就抛出异常
        if not pifile_id:
            errorinfo = {
                "FunctionType": change_info.get("functionType"),
                "TaskId": change_info.get("dmsTaskId"),
                "AiTaskId": task_id,
                "Reason": "未找到工艺信息表"
            }
            dmsoperator.post_errortodms(errorinfo=errorinfo)
            raise TaskException("未找到工艺信息表")
        if pifile_id:
            # 去掉工艺信息表中的不重要的信息
            pitext = self.cut_filetext(pifile_id)
            # 在剩余的内容中提取出需要的信息
            extract_list = [f"【{pi_query}】的处方", f"【{pi_query}】的生产工艺、设备信息、关键中间体的质量标准、关键工序及控制参数"]
            for extract in extract_list:
                extract_query = self.prompt.extract_everyinfo_prompt(query=extract, article=pitext)
                extract_onetext = gen_chain(extract_query)
                # print("*"*100)
                # print(f"{extract}的内容为：\n{extract_onetext}")
                # print("*"*100)
                # 输出提取的信息到前端
                pub_gen(task_id, extract_onetext, SOPfile.get("docId"), SOPfile.get("taskDocId"))
                extract_text += extract_onetext
        # 解析模板工艺规程文件
        print(extract_text)
        file_pubinfo, parainfo, tableinfo = self.parse_SOPjsonfile(SOPfile)

        gen_filename = "生成的文件"
        table_i = -1
        # 修改段落
        for i, para in enumerate(parainfo):
            # if not template_infos[i]["source"]:
            #     revise_SOPtext_dict = self.gen_SOPfile_para_table(para, extract_text, task_id, SOPfile)
            # else:
            #     revise_SOPtext_dict = self.gen_SOPfile_para_table_temp(para, extract_text, template_infos[i]["source"], task_id, SOPfile)
            revise_SOPtext_dict = self.gen_SOPfile_para_table(para, extract_text, task_id, SOPfile)

            edit_para_dict = {
                    "index": revise_SOPtext_dict["index"],
                    "sectionIndex": revise_SOPtext_dict["sectionIndex"],
                    "bodyChildIndex": revise_SOPtext_dict["bodyChildIndex"],
                    "operationType": 1,
                    "paragraphInfo":{
                        "paragraphChildInfos": [],
                        "referenceParagraphType": 2
                    }
                }
            # 插入纯文本和化学式下标
            textinfo_list = split_sub_string(revise_SOPtext_dict["text"])
            textinfos = []
            for textinfo_index, textinfo in enumerate(textinfo_list):
                textinfo_dict = {
                        "paragraphChildIndex": textinfo_index,
                        "paraChildType": 0,
                        "content": textinfo.get("content"),
                        "subSuperScript": textinfo.get("subSuperScript")
                    }
                textinfos.append(textinfo_dict)
            if textinfos:
                edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
            # 插入数学式
            if revise_SOPtext_dict["latexMathInfos"]:
                latex_index = len(textinfos)
                mathinfos = []
                for latex_math_info in revise_SOPtext_dict["latexMathInfos"]:
                    mathinfo_dict = {
                        "paragraphChildIndex": latex_index,
                        "paraChildType": 2,
                        "content": latex_math_info["text"]
                    }
                    latex_index += 1
                    mathinfos.append(mathinfo_dict)
                edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += mathinfos
            file_pubinfo["editInfos"].append(edit_para_dict)
            # 删除段落
            dele_edit_dict = {
                "index": revise_SOPtext_dict["index"],
                "sectionIndex": revise_SOPtext_dict["sectionIndex"],
                "bodyChildIndex": revise_SOPtext_dict["bodyChildIndex"],
                "operationType": 2
            }
            file_pubinfo["editInfos"].append(dele_edit_dict)
            if i == 0:
                gen_filename = revise_SOPtext_dict["text"]
            table_i = i

        # 修改表格
        for table in tableinfo:
            # 以table为单位进行修改
            table_i += 1
            # if not template_infos[table_i]["source"]:
            #     revise_SOPtext_dict = self.gen_SOPfile_para_table(table, extract_text, task_id, SOPfile)
            # else:
            #     revise_SOPtext_dict = self.gen_SOPfile_para_table_temp(table, extract_text, template_infos[table_i]["source"], task_id, SOPfile)
            revise_SOPtext_dict = self.gen_SOPfile_para_table(table, extract_text, task_id, SOPfile)
            # 判断生成的表格行数和原表格行数是否一致
            # 一致，直接修改，不一致，删除表格，重新生成
            if len(revise_SOPtext_dict["tableCellInfos"]) == len(table["tableCellInfos"]):
                table_dict = {
                    "index": revise_SOPtext_dict["index"],
                    "sectionIndex": revise_SOPtext_dict["sectionIndex"],
                    "bodyChildIndex": revise_SOPtext_dict["bodyChildIndex"],
                    "operationType": 0,
                    "tableInfo": {
                        "cellEditInfos": []
                    }
                }
                for cell in revise_SOPtext_dict["tableCellInfos"]:
                    cell_dict = {
                        "rowIndex": cell["rowIndex"],
                        "columnIndex": cell["columnIndex"],
                        "cellChildInfos": []
                    }
                    for para in cell["paragraphInfos"]:
                        edit_para_dict = {
                                "index": para["index"],
                                "sectionIndex": para["sectionIndex"],
                                "bodyChildIndex": para["bodyChildIndex"],
                                "operationType": 1,
                                "paragraphInfo": {
                                    "paragraphChildInfos": [],
                                    "referenceParagraphType": 2
                                }
                            }
                        # 插入纯文本和下标
                        textinfo_list = split_sub_string(para["text"])
                        textinfos = []
                        for textinfo_index, textinfo in enumerate(textinfo_list):
                            textinfo_dict = {
                                "paragraphChildIndex": textinfo_index,
                                "paraChildType": 0,
                                "content": textinfo.get("content"),
                                "subSuperScript": textinfo.get("subSuperScript")
                            }
                            textinfos.append(textinfo_dict)
                        if textinfos:
                            edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
                        # 插入数学公式
                        if para["latexMathInfos"]:
                            latex_index = len(textinfos)
                            mathinfos = []
                            for latex_math_info in para["latexMathInfos"]:
                                mathinfo_dict = {
                                    "paragraphChildIndex": latex_index,
                                    "paraChildType": 2,
                                    "content": latex_math_info["text"]
                                }
                                latex_index += 1
                                mathinfos.append(mathinfo_dict)
                            edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += mathinfos
                        cell_dict["cellChildInfos"].append(edit_para_dict)
                        # 删除段落
                        dele_para_dict = {
                            "index": para["index"],
                            "sectionIndex": para["sectionIndex"],
                            "bodyChildIndex": para["bodyChildIndex"],
                            "operationType": 2
                        }
                        cell_dict["cellChildInfos"].append(dele_para_dict)
                    table_dict["tableInfo"]["cellEditInfos"].append(cell_dict)
                file_pubinfo["editInfos"].append(table_dict)
            # 不一致，删除表格，重新生成
            else:
                # 先插入表格
                # 计算表格行数和列数
                max_row = -1
                max_column = -1
                for cell in revise_SOPtext_dict["tableCellInfos"]:
                    max_row = max(cell["rowIndex"], max_row)
                    max_column = max(cell["columnIndex"], max_column)
                row_num = max_row + 1
                column_num = max_column + 1
                table_dict = {
                    "index": revise_SOPtext_dict["index"],
                    "sectionIndex": revise_SOPtext_dict["sectionIndex"],
                    "bodyChildIndex": revise_SOPtext_dict["bodyChildIndex"],
                    "operationType": 1,
                    "tableInfo": {
                        "tableInsertInfo": {
                            "rowCount": row_num,
                            "columnCount": column_num,
                            "cellInsertInfos": []
                        }
                    }
                }
                for cell in revise_SOPtext_dict["tableCellInfos"]:
                    cell_dict = {
                        "rowIndex": cell["rowIndex"],
                        "columnIndex": cell["columnIndex"],
                        "cellInsertContentInfos": []
                    }
                    for para in cell["paragraphInfos"]:

                        edit_para_dict = {
                                "index": para["index"],
                                "paragraphInfo": {
                                    "paragraphChildInfos": [],
                                }
                            }
                        # 先插入纯文本和下标
                        textinfo_list = split_sub_string(para["text"])
                        textinfos = []
                        for textinfo_index, textinfo in enumerate(textinfo_list):
                            textinfo_dict = {
                                "paragraphChildIndex": textinfo_index,
                                "paraChildType": 0,
                                "content": textinfo.get("content"),
                                "subSuperScript": textinfo.get("subSuperScript")
                            }
                            textinfos.append(textinfo_dict)
                        if textinfos:
                            edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
                        # 插入数学公式
                        if para["latexMathInfos"]:
                            latex_index = len(textinfos)
                            mathinfos = []
                            for latex_math_info in para["latexMathInfos"]:
                                mathinfo_dict = {
                                    "paragraphChildIndex": latex_index,
                                    "paraChildType": 2,
                                    "content": latex_math_info["text"]
                                }
                                latex_index += 1
                                mathinfos.append(mathinfo_dict)
                            edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += mathinfos
                        cell_dict["cellInsertContentInfos"].append(edit_para_dict)
                    table_dict["tableInfo"]["tableInsertInfo"]["cellInsertInfos"].append(cell_dict)
                file_pubinfo["editInfos"].append(table_dict)
                # 再删除表格
                dele_table = {
                    "index": revise_SOPtext_dict["index"],
                    "sectionIndex": revise_SOPtext_dict["sectionIndex"],
                    "bodyChildIndex": revise_SOPtext_dict["bodyChildIndex"],
                    "operationType": 2,
                }
                file_pubinfo["editInfos"].append(dele_table)


        print("-"*200)
        file_pubinfo_json = json.dumps(file_pubinfo, indent=4, ensure_ascii=False)
        print(file_pubinfo_json)
        # # 保存修改的json文件
        # edit_SOPjsonfile_path = f"修改.json"
        # with open(edit_SOPjsonfile_path, "w", encoding="utf-8") as f:
        #     f.write(file_pubinfo_json)
        # print("-"*200)
        # 修改文件
        gen_res = dmsoperator.post_edit_file_v2(file_pubinfo)
        return gen_res["data"], gen_filename+".docx"




